{
 "cells": [
  {
   "cell_type": "code",
   "id": "initial_id",
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2025-09-17T04:48:25.628688Z",
     "start_time": "2025-09-17T04:48:22.451636Z"
    }
   },
   "source": [
    "import torch\n",
    "import torchvision\n",
    "from torch import nn\n",
    "from torch.nn import Sequential, Conv2d, MaxPool2d, Flatten\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "from torchvision import transforms\n",
    "\n",
    "import time\n",
    "\n",
    "# 将图像数据进行归一化处理，使得模型更容易收敛\n",
    "transform = transforms.Compose([\n",
    "    transforms.Resize(224),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))\n",
    "])\n",
    "\n",
    "train_data = torchvision.datasets.CIFAR10(root='../dataset', train=True, download=True, transform=transform)\n",
    "train_loader = DataLoader(train_data, batch_size=128, shuffle=True)\n",
    "\n",
    "test_data = torchvision.datasets.CIFAR10(root='../dataset', train=False, download=True, transform=transform)\n",
    "test_loader = DataLoader(test_data, batch_size=128, shuffle=True)\n",
    "\n",
    "print(f\"训练集的长度为：{len(train_data)}\")\n",
    "print(f\"测试集的长度为：{len(test_data)}\")\n",
    "print(f\"特征尺度为:{train_data[0][0].shape}\")\n",
    "device = torch.device(\"mps\" if torch.backends.mps.is_available() else \"cpu\")"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "训练集的长度为：50000\n",
      "测试集的长度为：10000\n",
      "特征尺度为:torch.Size([3, 224, 224])\n"
     ]
    }
   ],
   "execution_count": 1
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-09-17T04:48:26.463845Z",
     "start_time": "2025-09-17T04:48:25.631646Z"
    }
   },
   "cell_type": "code",
   "source": [
    "drop = 0.4\n",
    "def vgg_block(num_convs, in_channels, out_channels):\n",
    "    layers = []\n",
    "    for _ in range(num_convs):\n",
    "        layers.append(Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=3, padding=1))\n",
    "        layers.append(nn.ReLU(inplace=True))\n",
    "        in_channels = out_channels\n",
    "    layers.append(MaxPool2d(kernel_size=2, stride=2))\n",
    "    return nn.Sequential(*layers)\n",
    "\n",
    "class vgg(nn.Module):\n",
    "    def __init__(self, conv_arch, in_channels, inputs, outputs):\n",
    "        super(vgg, self).__init__()\n",
    "        conv_blocks = []\n",
    "        for (num_convs, out_channels) in conv_arch:\n",
    "            conv_blocks.append(vgg_block(num_convs, in_channels, out_channels))\n",
    "            in_channels = out_channels\n",
    "        size = inputs // 2 ** len(conv_arch)\n",
    "        self.model = nn.Sequential(\n",
    "            *conv_blocks,\n",
    "            nn.Flatten(),\n",
    "            nn.Linear(in_channels * size * size, 4096),\n",
    "            nn.Linear(4096, 4096),\n",
    "            nn.Linear(4096, outputs)\n",
    "        )\n",
    "\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.model(x)\n",
    "        return x\n",
    "conv_arch = ((1, 64), (1, 128), (2, 256), (2, 512), (2, 512))\n",
    "in_channels = 3\n",
    "inputs = 224\n",
    "outputs = 10\n",
    "model = vgg(conv_arch, in_channels=in_channels, inputs=inputs, outputs=outputs).to(device)\n",
    "\n",
    "from torchinfo import summary\n",
    "print(summary(model, (1, 3, 224, 224)))\n"
   ],
   "id": "8a58eb62b7dfdc72",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==========================================================================================\n",
      "Layer (type:depth-idx)                   Output Shape              Param #\n",
      "==========================================================================================\n",
      "vgg                                      [1, 10]                   --\n",
      "├─Sequential: 1-1                        [1, 10]                   --\n",
      "│    └─Sequential: 2-1                   [1, 64, 112, 112]         --\n",
      "│    │    └─Conv2d: 3-1                  [1, 64, 224, 224]         1,792\n",
      "│    │    └─ReLU: 3-2                    [1, 64, 224, 224]         --\n",
      "│    │    └─MaxPool2d: 3-3               [1, 64, 112, 112]         --\n",
      "│    └─Sequential: 2-2                   [1, 128, 56, 56]          --\n",
      "│    │    └─Conv2d: 3-4                  [1, 128, 112, 112]        73,856\n",
      "│    │    └─ReLU: 3-5                    [1, 128, 112, 112]        --\n",
      "│    │    └─MaxPool2d: 3-6               [1, 128, 56, 56]          --\n",
      "│    └─Sequential: 2-3                   [1, 256, 28, 28]          --\n",
      "│    │    └─Conv2d: 3-7                  [1, 256, 56, 56]          295,168\n",
      "│    │    └─ReLU: 3-8                    [1, 256, 56, 56]          --\n",
      "│    │    └─Conv2d: 3-9                  [1, 256, 56, 56]          590,080\n",
      "│    │    └─ReLU: 3-10                   [1, 256, 56, 56]          --\n",
      "│    │    └─MaxPool2d: 3-11              [1, 256, 28, 28]          --\n",
      "│    └─Sequential: 2-4                   [1, 512, 14, 14]          --\n",
      "│    │    └─Conv2d: 3-12                 [1, 512, 28, 28]          1,180,160\n",
      "│    │    └─ReLU: 3-13                   [1, 512, 28, 28]          --\n",
      "│    │    └─Conv2d: 3-14                 [1, 512, 28, 28]          2,359,808\n",
      "│    │    └─ReLU: 3-15                   [1, 512, 28, 28]          --\n",
      "│    │    └─MaxPool2d: 3-16              [1, 512, 14, 14]          --\n",
      "│    └─Sequential: 2-5                   [1, 512, 7, 7]            --\n",
      "│    │    └─Conv2d: 3-17                 [1, 512, 14, 14]          2,359,808\n",
      "│    │    └─ReLU: 3-18                   [1, 512, 14, 14]          --\n",
      "│    │    └─Conv2d: 3-19                 [1, 512, 14, 14]          2,359,808\n",
      "│    │    └─ReLU: 3-20                   [1, 512, 14, 14]          --\n",
      "│    │    └─MaxPool2d: 3-21              [1, 512, 7, 7]            --\n",
      "│    └─Flatten: 2-6                      [1, 25088]                --\n",
      "│    └─Linear: 2-7                       [1, 4096]                 102,764,544\n",
      "│    └─Linear: 2-8                       [1, 4096]                 16,781,312\n",
      "│    └─Linear: 2-9                       [1, 10]                   40,970\n",
      "==========================================================================================\n",
      "Total params: 128,807,306\n",
      "Trainable params: 128,807,306\n",
      "Non-trainable params: 0\n",
      "Total mult-adds (G): 7.61\n",
      "==========================================================================================\n",
      "Input size (MB): 0.60\n",
      "Forward/backward pass size (MB): 59.47\n",
      "Params size (MB): 515.23\n",
      "Estimated Total Size (MB): 575.31\n",
      "==========================================================================================\n"
     ]
    }
   ],
   "execution_count": 2
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-09-17T04:48:26.602027Z",
     "start_time": "2025-09-17T04:48:26.537858Z"
    }
   },
   "cell_type": "code",
   "source": [
    "X = torch.randn(1,3, 224, 224)\n",
    "for layer in model.model:\n",
    "    X = layer(X)\n",
    "    print(layer.__class__.__name__, f'output shape:{X.shape}')"
   ],
   "id": "9f62b2b71da0f767",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sequential output shape:torch.Size([1, 64, 112, 112])\n",
      "Sequential output shape:torch.Size([1, 128, 56, 56])\n",
      "Sequential output shape:torch.Size([1, 256, 28, 28])\n",
      "Sequential output shape:torch.Size([1, 512, 14, 14])\n",
      "Sequential output shape:torch.Size([1, 512, 7, 7])\n",
      "Flatten output shape:torch.Size([1, 25088])\n",
      "Linear output shape:torch.Size([1, 4096])\n",
      "Linear output shape:torch.Size([1, 4096])\n",
      "Linear output shape:torch.Size([1, 10])\n"
     ]
    }
   ],
   "execution_count": 3
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-09-17T04:48:26.613118Z",
     "start_time": "2025-09-17T04:48:26.609424Z"
    }
   },
   "cell_type": "code",
   "source": [
    "loss_fn = nn.CrossEntropyLoss()\n",
    "if torch.backends.mps.is_available():\n",
    "    loss_fn = loss_fn.to(device)\n",
    "learning_rate = 1e-4\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)\n",
    "epochs = 10\n",
    "train_step = 0\n",
    "\n",
    "\n",
    "# tensorboard记录一下训练时的 loss 变换\n",
    "writer = SummaryWriter(\"./logs-train/vgg\")"
   ],
   "id": "c4b0d7854144c2df",
   "outputs": [],
   "execution_count": 4
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-09-17T04:48:58.867715Z",
     "start_time": "2025-09-17T04:48:26.632307Z"
    }
   },
   "cell_type": "code",
   "source": [
    "from d2l_learn.utils import *\n",
    "train_CIAFAR10(model, train_data, train_loader, test_data, test_loader, loss_fn, optimizer, epochs, device, writer)"
   ],
   "id": "633293b19c78b323",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-------------第0轮训练开始了------------\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001B[0;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[0;31mKeyboardInterrupt\u001B[0m                         Traceback (most recent call last)",
      "Cell \u001B[0;32mIn[5], line 2\u001B[0m\n\u001B[1;32m      1\u001B[0m \u001B[38;5;28;01mfrom\u001B[39;00m\u001B[38;5;250m \u001B[39m\u001B[38;5;21;01md2l_learn\u001B[39;00m\u001B[38;5;21;01m.\u001B[39;00m\u001B[38;5;21;01mutils\u001B[39;00m\u001B[38;5;250m \u001B[39m\u001B[38;5;28;01mimport\u001B[39;00m \u001B[38;5;241m*\u001B[39m\n\u001B[0;32m----> 2\u001B[0m \u001B[43mtrain_CIAFAR10\u001B[49m\u001B[43m(\u001B[49m\u001B[43mmodel\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mtrain_data\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mtrain_loader\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mtest_data\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mtest_loader\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mloss_fn\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43moptimizer\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mepochs\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mdevice\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mwriter\u001B[49m\u001B[43m)\u001B[49m\n",
      "File \u001B[0;32m~/Documents/pythonProjects/python_study/d2l_learn/utils/__init__.py:203\u001B[0m, in \u001B[0;36mtrain_CIAFAR10\u001B[0;34m(model, train_data, train_loader, test_data, test_loader, loss_fn, optimizer, epochs, device, writer)\u001B[0m\n\u001B[1;32m    201\u001B[0m images, labels \u001B[38;5;241m=\u001B[39m data\n\u001B[1;32m    202\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m torch\u001B[38;5;241m.\u001B[39mbackends\u001B[38;5;241m.\u001B[39mmps\u001B[38;5;241m.\u001B[39mis_available():\n\u001B[0;32m--> 203\u001B[0m     images \u001B[38;5;241m=\u001B[39m \u001B[43mimages\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mto\u001B[49m\u001B[43m(\u001B[49m\u001B[43mdevice\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m    204\u001B[0m     labels \u001B[38;5;241m=\u001B[39m labels\u001B[38;5;241m.\u001B[39mto(device)\n\u001B[1;32m    205\u001B[0m out_labels \u001B[38;5;241m=\u001B[39m model(images)\n",
      "\u001B[0;31mKeyboardInterrupt\u001B[0m: "
     ]
    }
   ],
   "execution_count": 5
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
